# 加载取数与绘图所需的函数包
import datetime
import os

import pandas as pd
from hs_udata import set_token, stock_quote_daily


def GetData(stock_code, var1, var2):
    nowDate = datetime.datetime.now().strftime("%Y-%m-%d")

    date = datetime.datetime.strptime(nowDate, "%Y-%m-%d")
    date = date - datetime.timedelta(days=var2)  # 设置var2天前的日期

    data = pd.DataFrame([])
    set_token(token='ruXPpxAnX4-5DOF7SXgKOuI4y7WOtBPVXqE79Nj2q30dHluNv6Hd-ylE-JwynvBZ')  # 注册后，获取并替换token
    while len(data) <= var1:

        # 获取日行情数据，接口说明见 https://udata.hs.net/datas/332/
        # adjust_way枚举值为：0-不复权，1-前复权，2-后复权，此处取前复权
        data_i = stock_quote_daily(en_prod_code=stock_code, trading_date=date.strftime('%Y%m%d'), adjust_way=1)
        if data_i.turnover_status.values == '交易':
            data_i.set_index("prod_code", inplace=True)
            data = pd.concat([data_i, data], axis=0)  # 将行情数据按行拼接
        print(date, len(data))
        date += datetime.timedelta(days=-1)  # 日期变量自增

    # 返回行情数据
    return data


def getOldData(code, var1, var2):  # 获取var1条数  从距离现在var2天开始
    if os.path.exists("test.csv"):
        data = pd.read_csv("test.csv", parse_dates=True, encoding='UTF-8')
        if len(data.index) > var1:
            if code != data.prod_code.values[-1]:
                data = GetData(code, var1, var2)
                df = pd.DataFrame(data)
                df.to_csv('test.csv', encoding='UTF-8')  # 保存
        else:
            data = GetData(code, var1, var2)
            df = pd.DataFrame(data)
            df.to_csv('test.csv', encoding='UTF-8')  # 保存
    else:
        data = GetData(code, var1, var2)
        df = pd.DataFrame(data)
        df.to_csv('test.csv', encoding='UTF-8')  # 保存
        print(data)


getOldData("603232.SH", 200, 600)

